# 导入必要的库
import torch
from transformers import BertTokenizer, BertForNextSentencePrediction

# -------------------- 代码块 1：加载 BERT NSP 模型和分词器 --------------------
# 使用支持下一句预测任务的专用模型 (如果没有专门的 NSP 模型，可以继续使用 hfl/chinese-bert-wwm)
tokenizer = BertTokenizer.from_pretrained("hfl/chinese-bert-wwm")
model = BertForNextSentencePrediction.from_pretrained("hfl/chinese-bert-wwm")

# -------------------- 代码块 2：定义文本对 --------------------
sentence_1 = "今天我去超市买了些水果。"  # 句子 1 (修改为中文)
sentence_2 = "超市里的苹果和香蕉都很便宜。"  # 句子 2（符合上下文）
sentence_3 = "天气预报说明天会下雨。"  # 句子 3（不符合上下文）


# -------------------- 代码块 3：对文本对进行编码 --------------------
def encode_sentences(sent1, sent2, max_len=128):  # 添加 max_len 参数
    return tokenizer(
        text=sent1,
        text_pair=sent2,
        max_length=max_len,  # 使用参数
        padding='max_length',
        truncation=True,
        return_tensors="pt"
    )


encoding_positive = encode_sentences(sentence_1, sentence_2)  # 相关句子对
encoding_negative = encode_sentences(sentence_1, sentence_3)  # 无关句子对

# -------------------- 代码块 4：模型推理 --------------------
with torch.no_grad():
    # 正确的模型输出包含 next_sentence_logits
    output_positive = model(**encoding_positive)
    output_negative = model(**encoding_negative)

# 获取下一句预测的logits（形状为 [batch_size, 2]）
logits_positive = output_positive.logits  # logits[:, 0]: 不相关, logits[:, 1]: 相关
logits_negative = output_negative.logits

# 计算概率分布
probs_positive = torch.softmax(logits_positive, dim=1)
probs_negative = torch.softmax(logits_negative, dim=1)


# -------------------- 代码块 5：优化输出结果 --------------------
def format_probabilities(probs):
    return (f"不是上下句: {probs[0][0].item() * 100:.2f}%,"
            f" 是上下句: {probs[0][1].item() * 100:.2f}%")


print(f"句子1：{sentence_1}\n句子2：{sentence_2}")
print(f"预测结果：{format_probabilities(probs_positive)}")  # 显示两个类别的概率
print(f"预测类别：{'是上下句' if torch.argmax(probs_positive).item() == 1 else '不是上下句'}")

print("\n" + "-" * 50 + "\n")

print(f"句子1：{sentence_1}\n句子3：{sentence_3}")
print(f"预测结果：{format_probabilities(probs_negative)}")
print(f"预测类别：{'是上下句' if torch.argmax(probs_negative).item() == 1 else '不是上下句'}")


